{
 "cells": [
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:13.100063Z",
     "start_time": "2025-03-19T05:50:09.725107Z"
    }
   },
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import sklearn\n",
    "import pandas as pd\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "from tqdm.auto import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "print(sys.version_info)\n",
    "for module in mpl, np, pd, sklearn, torch:\n",
    "    print(module.__name__, module.__version__)\n",
    "    \n",
    "device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(device)\n",
    "\n",
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed_all(seed)\n",
    "np.random.seed(seed)\n"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sys.version_info(major=3, minor=12, micro=3, releaselevel='final', serial=0)\n",
      "matplotlib 3.10.0\n",
      "numpy 2.0.2\n",
      "pandas 2.2.3\n",
      "sklearn 1.6.1\n",
      "torch 2.6.0+cu126\n",
      "cuda:0\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 准备数据\n",
    "\n",
    "这里使用subword分词，我们使用已经清洗好的数据集，可以从[此处](https://www.kaggleusercontent.com/kf/98352223/eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0..nKt8lrIW5ej5QJQVpOWuqQ.oLIgiLMONU5Gpj_maVudRJa55NSOCILxk4JNZhvuXmeDBR-oG0uQm7bDHBfSwZRGvOBHQTsRV308iNP80btfwMinQ7yvJNt-GwdQF4XR4DIsg-2CbEPYiMsi_NdbL0FmE9LYStKdxCWbrCZCCMrTmo5LxR1txwibXaSpeP5Inobhbez5zetZIRH210CBuX2JbpRc_DULQpazKbtFPitwyfktVmdG_syvVAU6Sk9b0r0_erYAgb_jkKXX1Mxo1KzWSKLcAvbmMIPcsUkx9PmeJDs_wopfsQsZ1h5jaQX4_l0CTZrEenP6lIPDxpTwXANqqdHspmZeeEIAThqCHC6sb5DxTvG89BwzY9rc53Aa0uX4V806wJVybnRXoaV65K4GqpjnxbBK0WC8G-2lNtrqFujE89KDXZjFPgyfOEj1QIu13oFNSjgs6o4VV1PdZOrhiNdSmjb44c22l_unOaFojzJgzcPxq9AG2lcmrOpdZ2qu1jjdwey-58TA2ZHNCo3XnjEe2n3ignpnbsdLFpo22O8QakSUHv91wuYDYdNi3AiSmltL_k2ChuKfJ0G8kATpLe4k8wA26sO4GMXg4HImOr3b4aDVEIWXdApHP0ecFKs6ELTo8O7X-TK8Jvbua7e6qpDfDc-r_cD73fVSgSek5yNmKQMBzuVcjkprXmcxICQ.kV1b4N1s64NERKnt4zwQgQ/imdb_processed.csv)下载，分词使用 [subword-nmt](https://github.com/rsennrich/subword-nmt)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.660264Z",
     "start_time": "2025-03-19T05:50:23.407443Z"
    }
   },
   "source": [
    "cleaned_df = pd.read_csv(\"imdb_processed.csv\")\n",
    "print(cleaned_df.shape) # (50000, 2), 50000条评论, 2列\n",
    "\n",
    "# 随机打乱数据，取训练集和测试集\n",
    "np.random.seed(seed) #随机\n",
    "cleaned_df = cleaned_df.sample(frac=1).reset_index(drop=True)#打乱，frac=1表示全部打乱（frac是比例，reset_index(drop=True)是重新索引\n",
    "with open(\"imdb_train.txt\", \"w\", encoding=\"utf8\") as file:# 保存训练集\n",
    "    for line in cleaned_df.processed.values[:25000]:#只保存了processed列，即评论文本，没有保存label列\n",
    "        file.write(line.lower() + \"\\n\") #变为小写，token数量少一些\n",
    "\n",
    "with open(\"imdb_test.txt\", \"w\", encoding=\"utf8\") as file:# 保存测试集\n",
    "    for line in cleaned_df.processed.values[25000:]:\n",
    "        file.write(line.lower() + \"\\n\")"
   ],
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'imdb_processed.csv'",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mFileNotFoundError\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[2], line 1\u001B[0m\n\u001B[1;32m----> 1\u001B[0m cleaned_df \u001B[38;5;241m=\u001B[39m \u001B[43mpd\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mread_csv\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mimdb_processed.csv\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m)\u001B[49m\n\u001B[0;32m      2\u001B[0m \u001B[38;5;28mprint\u001B[39m(cleaned_df\u001B[38;5;241m.\u001B[39mshape) \u001B[38;5;66;03m# (50000, 2), 50000条评论, 2列\u001B[39;00m\n\u001B[0;32m      4\u001B[0m \u001B[38;5;66;03m# 随机打乱数据，取训练集和测试集\u001B[39;00m\n",
      "File \u001B[1;32mC:\\Program Files\\Python312\\Lib\\site-packages\\pandas\\io\\parsers\\readers.py:1026\u001B[0m, in \u001B[0;36mread_csv\u001B[1;34m(filepath_or_buffer, sep, delimiter, header, names, index_col, usecols, dtype, engine, converters, true_values, false_values, skipinitialspace, skiprows, skipfooter, nrows, na_values, keep_default_na, na_filter, verbose, skip_blank_lines, parse_dates, infer_datetime_format, keep_date_col, date_parser, date_format, dayfirst, cache_dates, iterator, chunksize, compression, thousands, decimal, lineterminator, quotechar, quoting, doublequote, escapechar, comment, encoding, encoding_errors, dialect, on_bad_lines, delim_whitespace, low_memory, memory_map, float_precision, storage_options, dtype_backend)\u001B[0m\n\u001B[0;32m   1013\u001B[0m kwds_defaults \u001B[38;5;241m=\u001B[39m _refine_defaults_read(\n\u001B[0;32m   1014\u001B[0m     dialect,\n\u001B[0;32m   1015\u001B[0m     delimiter,\n\u001B[1;32m   (...)\u001B[0m\n\u001B[0;32m   1022\u001B[0m     dtype_backend\u001B[38;5;241m=\u001B[39mdtype_backend,\n\u001B[0;32m   1023\u001B[0m )\n\u001B[0;32m   1024\u001B[0m kwds\u001B[38;5;241m.\u001B[39mupdate(kwds_defaults)\n\u001B[1;32m-> 1026\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43m_read\u001B[49m\u001B[43m(\u001B[49m\u001B[43mfilepath_or_buffer\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mkwds\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[1;32mC:\\Program Files\\Python312\\Lib\\site-packages\\pandas\\io\\parsers\\readers.py:620\u001B[0m, in \u001B[0;36m_read\u001B[1;34m(filepath_or_buffer, kwds)\u001B[0m\n\u001B[0;32m    617\u001B[0m _validate_names(kwds\u001B[38;5;241m.\u001B[39mget(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mnames\u001B[39m\u001B[38;5;124m\"\u001B[39m, \u001B[38;5;28;01mNone\u001B[39;00m))\n\u001B[0;32m    619\u001B[0m \u001B[38;5;66;03m# Create the parser.\u001B[39;00m\n\u001B[1;32m--> 620\u001B[0m parser \u001B[38;5;241m=\u001B[39m \u001B[43mTextFileReader\u001B[49m\u001B[43m(\u001B[49m\u001B[43mfilepath_or_buffer\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwds\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m    622\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m chunksize \u001B[38;5;129;01mor\u001B[39;00m iterator:\n\u001B[0;32m    623\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m parser\n",
      "File \u001B[1;32mC:\\Program Files\\Python312\\Lib\\site-packages\\pandas\\io\\parsers\\readers.py:1620\u001B[0m, in \u001B[0;36mTextFileReader.__init__\u001B[1;34m(self, f, engine, **kwds)\u001B[0m\n\u001B[0;32m   1617\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39moptions[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mhas_index_names\u001B[39m\u001B[38;5;124m\"\u001B[39m] \u001B[38;5;241m=\u001B[39m kwds[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mhas_index_names\u001B[39m\u001B[38;5;124m\"\u001B[39m]\n\u001B[0;32m   1619\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mhandles: IOHandles \u001B[38;5;241m|\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[1;32m-> 1620\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_engine \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_make_engine\u001B[49m\u001B[43m(\u001B[49m\u001B[43mf\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mengine\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[1;32mC:\\Program Files\\Python312\\Lib\\site-packages\\pandas\\io\\parsers\\readers.py:1880\u001B[0m, in \u001B[0;36mTextFileReader._make_engine\u001B[1;34m(self, f, engine)\u001B[0m\n\u001B[0;32m   1878\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mb\u001B[39m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;129;01min\u001B[39;00m mode:\n\u001B[0;32m   1879\u001B[0m         mode \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mb\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m-> 1880\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mhandles \u001B[38;5;241m=\u001B[39m \u001B[43mget_handle\u001B[49m\u001B[43m(\u001B[49m\n\u001B[0;32m   1881\u001B[0m \u001B[43m    \u001B[49m\u001B[43mf\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1882\u001B[0m \u001B[43m    \u001B[49m\u001B[43mmode\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1883\u001B[0m \u001B[43m    \u001B[49m\u001B[43mencoding\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43moptions\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mget\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mencoding\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43;01mNone\u001B[39;49;00m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1884\u001B[0m \u001B[43m    \u001B[49m\u001B[43mcompression\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43moptions\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mget\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mcompression\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43;01mNone\u001B[39;49;00m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1885\u001B[0m \u001B[43m    \u001B[49m\u001B[43mmemory_map\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43moptions\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mget\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mmemory_map\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43;01mFalse\u001B[39;49;00m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1886\u001B[0m \u001B[43m    \u001B[49m\u001B[43mis_text\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mis_text\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1887\u001B[0m \u001B[43m    \u001B[49m\u001B[43merrors\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43moptions\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mget\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mencoding_errors\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mstrict\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1888\u001B[0m \u001B[43m    \u001B[49m\u001B[43mstorage_options\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43moptions\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mget\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mstorage_options\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43;01mNone\u001B[39;49;00m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1889\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m   1890\u001B[0m \u001B[38;5;28;01massert\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mhandles \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[0;32m   1891\u001B[0m f \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mhandles\u001B[38;5;241m.\u001B[39mhandle\n",
      "File \u001B[1;32mC:\\Program Files\\Python312\\Lib\\site-packages\\pandas\\io\\common.py:873\u001B[0m, in \u001B[0;36mget_handle\u001B[1;34m(path_or_buf, mode, encoding, compression, memory_map, is_text, errors, storage_options)\u001B[0m\n\u001B[0;32m    868\u001B[0m \u001B[38;5;28;01melif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(handle, \u001B[38;5;28mstr\u001B[39m):\n\u001B[0;32m    869\u001B[0m     \u001B[38;5;66;03m# Check whether the filename is to be opened in binary mode.\u001B[39;00m\n\u001B[0;32m    870\u001B[0m     \u001B[38;5;66;03m# Binary mode does not support 'encoding' and 'newline'.\u001B[39;00m\n\u001B[0;32m    871\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m ioargs\u001B[38;5;241m.\u001B[39mencoding \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mb\u001B[39m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;129;01min\u001B[39;00m ioargs\u001B[38;5;241m.\u001B[39mmode:\n\u001B[0;32m    872\u001B[0m         \u001B[38;5;66;03m# Encoding\u001B[39;00m\n\u001B[1;32m--> 873\u001B[0m         handle \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mopen\u001B[39;49m\u001B[43m(\u001B[49m\n\u001B[0;32m    874\u001B[0m \u001B[43m            \u001B[49m\u001B[43mhandle\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    875\u001B[0m \u001B[43m            \u001B[49m\u001B[43mioargs\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mmode\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    876\u001B[0m \u001B[43m            \u001B[49m\u001B[43mencoding\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mioargs\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mencoding\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    877\u001B[0m \u001B[43m            \u001B[49m\u001B[43merrors\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43merrors\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    878\u001B[0m \u001B[43m            \u001B[49m\u001B[43mnewline\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m,\u001B[49m\n\u001B[0;32m    879\u001B[0m \u001B[43m        \u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m    880\u001B[0m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m    881\u001B[0m         \u001B[38;5;66;03m# Binary mode\u001B[39;00m\n\u001B[0;32m    882\u001B[0m         handle \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mopen\u001B[39m(handle, ioargs\u001B[38;5;241m.\u001B[39mmode)\n",
      "\u001B[1;31mFileNotFoundError\u001B[0m: [Errno 2] No such file or directory: 'imdb_processed.csv'"
     ]
    }
   ],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.661259Z",
     "start_time": "2025-03-19T05:50:23.661259Z"
    }
   },
   "source": "# !pip install subword-nmt  ",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# 学习bpe分词(很慢,学一次就好)\n",
    "# -i 选择学习的文件\n",
    "# -o 核心输出文件,分词需要用到imdb_bpe_code\n",
    "# --write-vocabulary 字典输出文件\n",
    "# -s 词表大小\n",
    "!subword-nmt learn-joint-bpe-and-vocab -i ./imdb_train.txt -o ./imdb_bpe_code --write-vocabulary ./imdb_bpe_vocab -s 8000\n",
    "\n",
    "\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "# 应用bpe分词,-c 指定 BPE 编码的配置文件\n",
    "!subword-nmt apply-bpe -c ./imdb_bpe_code -i ./imdb_train.txt -o ./imdb_train_bpe.txt\n",
    "!subword-nmt apply-bpe -c ./imdb_bpe_code -i ./imdb_test.txt -o ./imdb_test_bpe.txt"
   ],
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# 分词后的数据长什么样,与分词前imdb_train.txt进行对比来理解，@@ 是分词的标记，如果一个单词被分开，就会加上@@\n",
    "!head ./imdb_train_bpe.txt\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "cleaned_df.head()"
   ],
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "subwords = []\n",
    "with open(\"imdb_train_bpe.txt\", \"r\", encoding=\"utf8\") as file:\n",
    "    for line in file.readlines():\n",
    "        subwords.append(line.strip())\n",
    "        \n",
    "with open(\"imdb_test_bpe.txt\", \"r\", encoding=\"utf8\") as file:\n",
    "    for line in file.readlines():\n",
    "        subwords.append(line.strip())\n",
    "        \n",
    "cleaned_df[\"subwords10k\"] = subwords # 保存分词后的结果\n",
    "cleaned_df[\"split\"] = [\"train\"] * 25000 + [\"test\"] * 25000 # 标记训练集和测试集\n",
    "cleaned_df.to_csv(\"imdb_subwords.csv\", index=False) #把分词后的结果保存到csv文件"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "# 随后加载数据集就从bpe分词的文件里加载\n",
    "class IMDBDataset(Dataset):\n",
    "    def __init__(self, mode=\"train\"):\n",
    "        df = pd.read_csv(\"imdb_subwords.csv\").query(\"split == '{}'\".format(mode)) # 加载训练集或测试集，query语句筛选\n",
    "        self.texts = df[\"subwords10k\"].values # 评论文本\n",
    "        self.labels = df[\"label\"].values # 评论标签\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return self.texts[idx].split(), self.labels[idx] # 返回文本和标签\n",
    "    \n",
    "    \n",
    "train_ds = IMDBDataset(\"train\")\n",
    "test_ds = IMDBDataset(\"test\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "train_ds[0]"
   ],
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "execution_count": null
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构造 word2idx 和 idx2word"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "#载入词表，看下词表长度，词表就像英语字典\n",
    "word2idx = {\n",
    "    \"[PAD]\": 0,     # 填充 token\n",
    "    \"[BOS]\": 1,     # begin of sentence\n",
    "    \"[UNK]\": 2,     # 未知 token\n",
    "    \"[EOS]\": 3,     # end of sentence\n",
    "}\n",
    "idx2word = {value: key for key, value in word2idx.items()}\n",
    "print(idx2word)\n",
    "index = len(idx2word) # 词表长度，现在为4\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "threshold = 1  # 出现次数低于此的token舍弃\n",
    "with open(\"imdb_bpe_vocab\", \"r\", encoding=\"utf8\") as file:\n",
    "    for line in tqdm(file.readlines()):\n",
    "        token, counts = line.strip().split()\n",
    "        if int(counts) >= threshold:\n",
    "            word2idx[token] = index # 加入词表\n",
    "            idx2word[index] = token # 加入反向词典\n",
    "            index += 1\n",
    "\n",
    "vocab_size = len(word2idx)\n",
    "print(\"vocab_size: {}\".format(vocab_size))"
   ],
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.688845Z",
     "start_time": "2025-03-19T05:50:23.675263Z"
    }
   },
   "source": [
    "# 选择 max_length\n",
    "length_collect = {}\n",
    "for text, label in train_ds:\n",
    "    length = len(text)\n",
    "    length_collect[length] = length_collect.get(length, 0) + 1\n",
    "    \n",
    "MAX_LENGTH = 500\n",
    "plt.bar(length_collect.keys(), length_collect.values())\n",
    "plt.axvline(MAX_LENGTH, label=\"max length\", c=\"gray\", ls=\":\")\n",
    "plt.legend()\n",
    "plt.show()"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'train_ds' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[3], line 3\u001B[0m\n\u001B[0;32m      1\u001B[0m \u001B[38;5;66;03m# 选择 max_length\u001B[39;00m\n\u001B[0;32m      2\u001B[0m length_collect \u001B[38;5;241m=\u001B[39m {}\n\u001B[1;32m----> 3\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m text, label \u001B[38;5;129;01min\u001B[39;00m \u001B[43mtrain_ds\u001B[49m:\n\u001B[0;32m      4\u001B[0m     length \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mlen\u001B[39m(text)\n\u001B[0;32m      5\u001B[0m     length_collect[length] \u001B[38;5;241m=\u001B[39m length_collect\u001B[38;5;241m.\u001B[39mget(length, \u001B[38;5;241m0\u001B[39m) \u001B[38;5;241m+\u001B[39m \u001B[38;5;241m1\u001B[39m\n",
      "\u001B[1;31mNameError\u001B[0m: name 'train_ds' is not defined"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.716342Z",
     "start_time": "2025-03-19T05:50:23.695843Z"
    }
   },
   "source": [
    "class Tokenizer:\n",
    "    def __init__(self, word2idx, idx2word, max_length=500, pad_idx=0, bos_idx=1, eos_idx=3, unk_idx=2):\n",
    "        self.word2idx = word2idx\n",
    "        self.idx2word = idx2word\n",
    "        self.max_length = max_length\n",
    "        self.pad_idx = pad_idx\n",
    "        self.bos_idx = bos_idx\n",
    "        self.eos_idx = eos_idx\n",
    "        self.unk_idx = unk_idx\n",
    "    \n",
    "    def encode(self, text_list, padding_first=False):\n",
    "        \"\"\"如果padding_first == True，则padding加载前面，否则加载后面\"\"\"\n",
    "        max_length = min(self.max_length, 2 + max([len(text) for text in text_list]))\n",
    "        indices_list = []\n",
    "        for text in text_list:\n",
    "            indices = [self.bos_idx] + [self.word2idx.get(word, self.unk_idx) for word in text[:max_length-2]] + [self.eos_idx] #变为id，未登录词用unk_idx代替，句子前后加bos和eos\n",
    "            if padding_first: # padding加载前面\n",
    "                indices = [self.pad_idx] * (max_length - len(indices)) + indices\n",
    "            else:# padding加载后面\n",
    "                indices = indices + [self.pad_idx] * (max_length - len(indices))\n",
    "            indices_list.append(indices)\n",
    "        return torch.tensor(indices_list)\n",
    "    \n",
    "    \n",
    "    def decode(self, indices_list, remove_bos=True, remove_eos=True, remove_pad=True, split=False):\n",
    "        text_list = []\n",
    "        for indices in indices_list:\n",
    "            text = []\n",
    "            for index in indices:\n",
    "                word = self.idx2word.get(index, \"[UNK]\")\n",
    "                if remove_bos and word == \"[BOS]\":\n",
    "                    continue\n",
    "                if remove_eos and word == \"[EOS]\":\n",
    "                    break\n",
    "                if remove_pad and word == \"[PAD]\":\n",
    "                    break\n",
    "                text.append(word)\n",
    "            text_list.append(\" \".join(text) if not split else text)\n",
    "        return text_list\n",
    "    \n",
    "\n",
    "tokenizer = Tokenizer(word2idx=word2idx, idx2word=idx2word)\n",
    "raw_text = [\"hello world\".split(), \"i really liked sum@@ mer@@ sla@@ m due look a@@ ren@@ a , cur@@ \".split(), \"this is a test\".split()]\n",
    "indices = tokenizer.encode(raw_text, padding_first=False)\n",
    "\n",
    "print(\"raw text-------------------\")\n",
    "for raw in raw_text:\n",
    "    print(raw)\n",
    "print(\"indices---------------\")\n",
    "for index in indices:\n",
    "    print(index)\n"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'word2idx' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[4], line 42\u001B[0m\n\u001B[0;32m     38\u001B[0m             text_list\u001B[38;5;241m.\u001B[39mappend(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m \u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;241m.\u001B[39mjoin(text) \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m split \u001B[38;5;28;01melse\u001B[39;00m text)\n\u001B[0;32m     39\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m text_list\n\u001B[1;32m---> 42\u001B[0m tokenizer \u001B[38;5;241m=\u001B[39m Tokenizer(word2idx\u001B[38;5;241m=\u001B[39m\u001B[43mword2idx\u001B[49m, idx2word\u001B[38;5;241m=\u001B[39midx2word)\n\u001B[0;32m     43\u001B[0m raw_text \u001B[38;5;241m=\u001B[39m [\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mhello world\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;241m.\u001B[39msplit(), \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mi really liked sum@@ mer@@ sla@@ m due look a@@ ren@@ a , cur@@ \u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;241m.\u001B[39msplit(), \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mthis is a test\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;241m.\u001B[39msplit()]\n\u001B[0;32m     44\u001B[0m indices \u001B[38;5;241m=\u001B[39m tokenizer\u001B[38;5;241m.\u001B[39mencode(raw_text, padding_first\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m)\n",
      "\u001B[1;31mNameError\u001B[0m: name 'word2idx' is not defined"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.826219Z",
     "start_time": "2025-03-19T05:50:23.815360Z"
    }
   },
   "cell_type": "code",
   "source": "idx2word[2230]",
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'idx2word' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[5], line 1\u001B[0m\n\u001B[1;32m----> 1\u001B[0m \u001B[43midx2word\u001B[49m[\u001B[38;5;241m2230\u001B[39m]\n",
      "\u001B[1;31mNameError\u001B[0m: name 'idx2word' is not defined"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.878457Z",
     "start_time": "2025-03-19T05:50:23.864223Z"
    }
   },
   "cell_type": "code",
   "source": [
    "decode_text = tokenizer.decode(indices.tolist(), remove_bos=False, remove_eos=False, remove_pad=False)\n",
    "print(\"decode text\")\n",
    "for decode in decode_text:\n",
    "    print(decode)"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'tokenizer' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[6], line 1\u001B[0m\n\u001B[1;32m----> 1\u001B[0m decode_text \u001B[38;5;241m=\u001B[39m \u001B[43mtokenizer\u001B[49m\u001B[38;5;241m.\u001B[39mdecode(indices\u001B[38;5;241m.\u001B[39mtolist(), remove_bos\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m, remove_eos\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m, remove_pad\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m)\n\u001B[0;32m      2\u001B[0m \u001B[38;5;28mprint\u001B[39m(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mdecode text\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[0;32m      3\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m decode \u001B[38;5;129;01min\u001B[39;00m decode_text:\n",
      "\u001B[1;31mNameError\u001B[0m: name 'tokenizer' is not defined"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.934706Z",
     "start_time": "2025-03-19T05:50:23.931008Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import re\n",
    "\n",
    "# 输入字符串\n",
    "text = \"i really liked sum@@ mer@@ sla@@ m due look a@@ ren@@ a , cur@@ \"\n",
    "\n",
    "# 使用正则表达式替换 \"@@ \" 为空字符串\n",
    "cleaned_text = re.sub(r'@@\\s*', '', text)\n",
    "\n",
    "print(cleaned_text)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i really liked summerslam due look arena , cur\n"
     ]
    }
   ],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:23.981246Z",
     "start_time": "2025-03-19T05:50:23.965710Z"
    }
   },
   "source": [
    "def collate_fct(batch):\n",
    "    \"\"\"\n",
    "    把字符串列表转化为tensor\n",
    "    :param batch:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    text_list = [item[0] for item in batch] #batch中每个item的第一个元素是text,是输入，类型为list\n",
    "    label_list = [item[1] for item in batch] #batch中每个item的第二个元素是label,是输出，类型为int\n",
    "    # 这里使用 padding first\n",
    "    text_list = tokenizer.encode(text_list, padding_first=True).to(dtype=torch.int)\n",
    "    return text_list, torch.tensor(label_list).reshape(-1, 1).to(dtype=torch.float)\n",
    "\n",
    "batch_size = 128\n",
    "train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True, collate_fn=collate_fct)\n",
    "test_dl = DataLoader(test_ds, batch_size=batch_size, shuffle=False, collate_fn=collate_fct)"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'DataLoader' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[8], line 14\u001B[0m\n\u001B[0;32m     11\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m text_list, torch\u001B[38;5;241m.\u001B[39mtensor(label_list)\u001B[38;5;241m.\u001B[39mreshape(\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m, \u001B[38;5;241m1\u001B[39m)\u001B[38;5;241m.\u001B[39mto(dtype\u001B[38;5;241m=\u001B[39mtorch\u001B[38;5;241m.\u001B[39mfloat)\n\u001B[0;32m     13\u001B[0m batch_size \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m128\u001B[39m\n\u001B[1;32m---> 14\u001B[0m train_dl \u001B[38;5;241m=\u001B[39m \u001B[43mDataLoader\u001B[49m(train_ds, batch_size\u001B[38;5;241m=\u001B[39mbatch_size, shuffle\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mTrue\u001B[39;00m, collate_fn\u001B[38;5;241m=\u001B[39mcollate_fct)\n\u001B[0;32m     15\u001B[0m test_dl \u001B[38;5;241m=\u001B[39m DataLoader(test_ds, batch_size\u001B[38;5;241m=\u001B[39mbatch_size, shuffle\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m, collate_fn\u001B[38;5;241m=\u001B[39mcollate_fct)\n",
      "\u001B[1;31mNameError\u001B[0m: name 'DataLoader' is not defined"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 定义模型"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:24.028385Z",
     "start_time": "2025-03-19T05:50:24.003248Z"
    }
   },
   "source": [
    "class LSTM(nn.Module):\n",
    "    def __init__(self, embedding_dim=16, hidden_dim=64, vocab_size=vocab_size, num_layers=1, bidirectional=False):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.embeding = nn.Embedding(vocab_size, embedding_dim)\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=num_layers, batch_first=True, bidirectional=bidirectional)\n",
    "        self.layer = nn.Linear(hidden_dim * (2 if bidirectional else 1), hidden_dim)\n",
    "        self.fc = nn.Linear(hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # [bs, seq length]\n",
    "        x = self.embeding(x)\n",
    "        # [bs, seq length, embedding_dim] -> shape [bs, seq length, hidden_dim ]\n",
    "        seq_output, (hidden, cell) = self.lstm(x)\n",
    "\n",
    "        x = seq_output[:, -1, :]\n",
    "        # 取最后一个时间步的输出 (这也是为什么要设置padding_first=True的原因)\n",
    "        x = self.layer(x)\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "    \n",
    "sample_inputs = torch.randint(0, vocab_size, (2, 128))\n",
    "    \n",
    "print(\"{:=^80}\".format(\" 一层单向 LSTM \"))       \n",
    "for key, value in LSTM().named_parameters():\n",
    "    print(f\"{key:^40}paramerters num: {np.prod(value.shape)}\")\n",
    "\n",
    "    \n",
    "print(\"{:=^80}\".format(\" 一层双向 LSTM \"))       \n",
    "for key, value in LSTM(bidirectional=True).named_parameters():\n",
    "    print(f\"{key:^40}paramerters num: {np.prod(value.shape)}\")\n",
    "\n",
    "    \n",
    "print(\"{:=^80}\".format(\" 两层单向 LSTM \"))       \n",
    "for key, value in LSTM(num_layers=2).named_parameters():\n",
    "    print(f\"{key:^40}paramerters num: {np.prod(value.shape)}\")\n"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'vocab_size' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[9], line 1\u001B[0m\n\u001B[1;32m----> 1\u001B[0m \u001B[38;5;28;43;01mclass\u001B[39;49;00m\u001B[43m \u001B[49m\u001B[38;5;21;43;01mLSTM\u001B[39;49;00m\u001B[43m(\u001B[49m\u001B[43mnn\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mModule\u001B[49m\u001B[43m)\u001B[49m\u001B[43m:\u001B[49m\n\u001B[0;32m      2\u001B[0m \u001B[43m    \u001B[49m\u001B[38;5;28;43;01mdef\u001B[39;49;00m\u001B[43m \u001B[49m\u001B[38;5;21;43m__init__\u001B[39;49m\u001B[43m(\u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43membedding_dim\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;241;43m16\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mhidden_dim\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;241;43m64\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mvocab_size\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mvocab_size\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mnum_layers\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;241;43m1\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mbidirectional\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mFalse\u001B[39;49;00m\u001B[43m)\u001B[49m\u001B[43m:\u001B[49m\n\u001B[0;32m      3\u001B[0m \u001B[43m        \u001B[49m\u001B[38;5;28;43msuper\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mLSTM\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m)\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[38;5;21;43m__init__\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n",
      "Cell \u001B[1;32mIn[9], line 2\u001B[0m, in \u001B[0;36mLSTM\u001B[1;34m()\u001B[0m\n\u001B[0;32m      1\u001B[0m \u001B[38;5;28;01mclass\u001B[39;00m \u001B[38;5;21;01mLSTM\u001B[39;00m(nn\u001B[38;5;241m.\u001B[39mModule):\n\u001B[1;32m----> 2\u001B[0m     \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m__init__\u001B[39m(\u001B[38;5;28mself\u001B[39m, embedding_dim\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m16\u001B[39m, hidden_dim\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m64\u001B[39m, vocab_size\u001B[38;5;241m=\u001B[39m\u001B[43mvocab_size\u001B[49m, num_layers\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m1\u001B[39m, bidirectional\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m):\n\u001B[0;32m      3\u001B[0m         \u001B[38;5;28msuper\u001B[39m(LSTM, \u001B[38;5;28mself\u001B[39m)\u001B[38;5;241m.\u001B[39m\u001B[38;5;21m__init__\u001B[39m()\n\u001B[0;32m      4\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39membeding \u001B[38;5;241m=\u001B[39m nn\u001B[38;5;241m.\u001B[39mEmbedding(vocab_size, embedding_dim)\n",
      "\u001B[1;31mNameError\u001B[0m: name 'vocab_size' is not defined"
     ]
    }
   ],
   "execution_count": 9
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:24.029372Z",
     "start_time": "2025-03-19T05:50:24.029372Z"
    }
   },
   "source": [
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "@torch.no_grad()\n",
    "def evaluating(model, dataloader, loss_fct):\n",
    "    loss_list = []\n",
    "    pred_list = []\n",
    "    label_list = []\n",
    "    for datas, labels in dataloader:\n",
    "        datas = datas.to(device)\n",
    "        labels = labels.to(device)\n",
    "        # 前向计算\n",
    "        logits = model(datas)\n",
    "        loss = loss_fct(logits, labels)         # 验证集损失\n",
    "        loss_list.append(loss.item())\n",
    "        # 二分类\n",
    "        preds = logits > 0\n",
    "        pred_list.extend(preds.cpu().numpy().tolist())\n",
    "        label_list.extend(labels.cpu().numpy().tolist())\n",
    "        \n",
    "    acc = accuracy_score(label_list, pred_list)\n",
    "    return np.mean(loss_list), acc\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TensorBoard 可视化\n",
    "\n",
    "\n",
    "训练过程中可以使用如下命令启动tensorboard服务。\n",
    "\n",
    "```shell\n",
    "tensorboard \\\n",
    "    --logdir=runs \\     # log 存放路径\n",
    "    --host 0.0.0.0 \\    # ip\n",
    "    --port 8848         # 端口\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:27.591Z",
     "start_time": "2025-03-19T05:50:24.032372Z"
    }
   },
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "\n",
    "class TensorBoardCallback:\n",
    "    def __init__(self, log_dir, flush_secs=10):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            log_dir (str): dir to write log.\n",
    "            flush_secs (int, optional): write to dsk each flush_secs seconds. Defaults to 10.\n",
    "        \"\"\"\n",
    "        self.writer = SummaryWriter(log_dir=log_dir, flush_secs=flush_secs)\n",
    "\n",
    "    def draw_model(self, model, input_shape):\n",
    "        self.writer.add_graph(model, input_to_model=torch.randn(input_shape))\n",
    "        \n",
    "    def add_loss_scalars(self, step, loss, val_loss):\n",
    "        self.writer.add_scalars(\n",
    "            main_tag=\"training/loss\", \n",
    "            tag_scalar_dict={\"loss\": loss, \"val_loss\": val_loss},\n",
    "            global_step=step,\n",
    "            )\n",
    "        \n",
    "    def add_acc_scalars(self, step, acc, val_acc):\n",
    "        self.writer.add_scalars(\n",
    "            main_tag=\"training/accuracy\",\n",
    "            tag_scalar_dict={\"accuracy\": acc, \"val_accuracy\": val_acc},\n",
    "            global_step=step,\n",
    "        )\n",
    "        \n",
    "    def add_lr_scalars(self, step, learning_rate):\n",
    "        self.writer.add_scalars(\n",
    "            main_tag=\"training/learning_rate\",\n",
    "            tag_scalar_dict={\"learning_rate\": learning_rate},\n",
    "            global_step=step,\n",
    "            \n",
    "        )\n",
    "    \n",
    "    def __call__(self, step, **kwargs):\n",
    "        # add loss\n",
    "        loss = kwargs.pop(\"loss\", None)\n",
    "        val_loss = kwargs.pop(\"val_loss\", None)\n",
    "        if loss is not None and val_loss is not None:\n",
    "            self.add_loss_scalars(step, loss, val_loss)\n",
    "        # add acc\n",
    "        acc = kwargs.pop(\"acc\", None)\n",
    "        val_acc = kwargs.pop(\"val_acc\", None)\n",
    "        if acc is not None and val_acc is not None:\n",
    "            self.add_acc_scalars(step, acc, val_acc)\n",
    "        # add lr\n",
    "        learning_rate = kwargs.pop(\"lr\", None)\n",
    "        if learning_rate is not None:\n",
    "            self.add_lr_scalars(step, learning_rate)\n"
   ],
   "outputs": [],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save Best\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:27.595899Z",
     "start_time": "2025-03-19T05:50:27.591510Z"
    }
   },
   "source": [
    "class SaveCheckpointsCallback:\n",
    "    def __init__(self, save_dir, save_step=5000, save_best_only=True):\n",
    "        \"\"\"\n",
    "        Save checkpoints each save_epoch epoch. \n",
    "        We save checkpoint by epoch in this implementation.\n",
    "        Usually, training scripts with pytorch evaluating model and save checkpoint by step.\n",
    "\n",
    "        Args:\n",
    "            save_dir (str): dir to save checkpoint\n",
    "            save_epoch (int, optional): the frequency to save checkpoint. Defaults to 1.\n",
    "            save_best_only (bool, optional): If True, only save the best model or save each model at every epoch.\n",
    "        \"\"\"\n",
    "        self.save_dir = save_dir\n",
    "        self.save_step = save_step\n",
    "        self.save_best_only = save_best_only\n",
    "        self.best_metrics = -1\n",
    "        \n",
    "        # mkdir\n",
    "        if not os.path.exists(self.save_dir):\n",
    "            os.mkdir(self.save_dir)\n",
    "        \n",
    "    def __call__(self, step, state_dict, metric=None):\n",
    "        if step % self.save_step > 0:\n",
    "            return\n",
    "        \n",
    "        if self.save_best_only:\n",
    "            assert metric is not None\n",
    "            if metric >= self.best_metrics:\n",
    "                # save checkpoints\n",
    "                torch.save(state_dict, os.path.join(self.save_dir, \"best.ckpt\"))\n",
    "                # update best metrics\n",
    "                self.best_metrics = metric\n",
    "        else:\n",
    "            torch.save(state_dict, os.path.join(self.save_dir, f\"{step}.ckpt\"))\n",
    "\n"
   ],
   "outputs": [],
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Early Stop"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:27.600076Z",
     "start_time": "2025-03-19T05:50:27.595899Z"
    }
   },
   "source": [
    "class EarlyStopCallback:\n",
    "    def __init__(self, patience=5, min_delta=0.01):\n",
    "        \"\"\"\n",
    "\n",
    "        Args:\n",
    "            patience (int, optional): Number of epochs with no improvement after which training will be stopped.. Defaults to 5.\n",
    "            min_delta (float, optional): Minimum change in the monitored quantity to qualify as an improvement, i.e. an absolute \n",
    "                change of less than min_delta, will count as no improvement. Defaults to 0.01.\n",
    "        \"\"\"\n",
    "        self.patience = patience\n",
    "        self.min_delta = min_delta\n",
    "        self.best_metric = -1\n",
    "        self.counter = 0\n",
    "        \n",
    "    def __call__(self, metric):\n",
    "        if metric >= self.best_metric + self.min_delta:\n",
    "            # update best metric\n",
    "            self.best_metric = metric\n",
    "            # reset counter \n",
    "            self.counter = 0\n",
    "        else: \n",
    "            self.counter += 1\n",
    "            \n",
    "    @property\n",
    "    def early_stop(self):\n",
    "        return self.counter >= self.patience\n"
   ],
   "outputs": [],
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:28.106053Z",
     "start_time": "2025-03-19T05:50:27.601073Z"
    }
   },
   "source": [
    "# 训练\n",
    "def training(\n",
    "    model, \n",
    "    train_loader, \n",
    "    val_loader, \n",
    "    epoch, \n",
    "    loss_fct, \n",
    "    optimizer, \n",
    "    tensorboard_callback=None,\n",
    "    save_ckpt_callback=None,\n",
    "    early_stop_callback=None,\n",
    "    eval_step=500,\n",
    "    ):\n",
    "    record_dict = {\n",
    "        \"train\": [],\n",
    "        \"val\": []\n",
    "    }\n",
    "    \n",
    "    global_step = 0\n",
    "    model.train()\n",
    "    with tqdm(total=epoch * len(train_loader)) as pbar:\n",
    "        for epoch_id in range(epoch):\n",
    "            # training\n",
    "            for datas, labels in train_loader:\n",
    "                datas = datas.to(device)\n",
    "                labels = labels.to(device)\n",
    "                # 梯度清空\n",
    "                optimizer.zero_grad()\n",
    "                # 模型前向计算\n",
    "                logits = model(datas)\n",
    "                # 计算损失\n",
    "                loss = loss_fct(logits, labels)\n",
    "                # 梯度回传\n",
    "                loss.backward()\n",
    "                # 调整优化器，包括学习率的变动等\n",
    "                optimizer.step()\n",
    "                preds = logits > 0\n",
    "            \n",
    "                acc = accuracy_score(labels.cpu().numpy(), preds.cpu().numpy())    \n",
    "                loss = loss.cpu().item()\n",
    "                # record\n",
    "                \n",
    "                record_dict[\"train\"].append({\n",
    "                    \"loss\": loss, \"acc\": acc, \"step\": global_step\n",
    "                })\n",
    "                \n",
    "                # evaluating\n",
    "                if global_step % eval_step == 0:\n",
    "                    model.eval()\n",
    "                    val_loss, val_acc = evaluating(model, val_loader, loss_fct)\n",
    "                    record_dict[\"val\"].append({\n",
    "                        \"loss\": val_loss, \"acc\": val_acc, \"step\": global_step\n",
    "                    })\n",
    "                    model.train()\n",
    "                    \n",
    "                    # 1. 使用 tensorboard 可视化\n",
    "                    if tensorboard_callback is not None:\n",
    "                        tensorboard_callback(\n",
    "                            global_step, \n",
    "                            loss=loss, val_loss=val_loss,\n",
    "                            acc=acc, val_acc=val_acc,\n",
    "                            lr=optimizer.param_groups[0][\"lr\"],\n",
    "                            )\n",
    "                \n",
    "                    # 2. 保存模型权重 save model checkpoint\n",
    "                    if save_ckpt_callback is not None:\n",
    "                        save_ckpt_callback(global_step, model.state_dict(), metric=val_acc)\n",
    "\n",
    "                    # 3. 早停 Early Stop\n",
    "                    if early_stop_callback is not None:\n",
    "                        early_stop_callback(val_acc)\n",
    "                        if early_stop_callback.early_stop:\n",
    "                            print(f\"Early stop at epoch {epoch_id} / global_step {global_step}\")\n",
    "                            return record_dict\n",
    "                    \n",
    "                # udate step\n",
    "                global_step += 1\n",
    "                pbar.update(1)\n",
    "                pbar.set_postfix({\"epoch\": epoch_id})\n",
    "        \n",
    "    return record_dict\n",
    "        \n",
    "\n",
    "epoch = 20\n",
    "\n",
    "model = LSTM()\n",
    "\n",
    "# 1. 定义损失函数 采用交叉熵损失 (但是二分类)\n",
    "loss_fct = F.binary_cross_entropy_with_logits\n",
    "# 2. 定义优化器 采用 adam\n",
    "# Optimizers specified in the torch.optim package\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# 1. tensorboard 可视化\n",
    "if not os.path.exists(\"runs\"):\n",
    "    os.mkdir(\"runs\")\n",
    "tensorboard_callback = TensorBoardCallback(\"runs/imdb-lstm-subword\")\n",
    "# tensorboard_callback.draw_model(model, [1, MAX_LENGTH])\n",
    "# 2. save best\n",
    "if not os.path.exists(\"checkpoints\"):\n",
    "    os.makedirs(\"checkpoints\")\n",
    "save_ckpt_callback = SaveCheckpointsCallback(\"checkpoints/imdb-lstm-subword\", save_step=len(train_dl), save_best_only=True)\n",
    "# 3. early stop\n",
    "early_stop_callback = EarlyStopCallback(patience=10)\n",
    "\n",
    "model = model.to(device)\n",
    "record = training(\n",
    "    model, \n",
    "    train_dl, \n",
    "    test_dl, \n",
    "    epoch, \n",
    "    loss_fct, \n",
    "    optimizer, \n",
    "    tensorboard_callback=tensorboard_callback,\n",
    "    save_ckpt_callback=save_ckpt_callback,\n",
    "    early_stop_callback=early_stop_callback,\n",
    "    eval_step=len(train_dl)\n",
    "    )"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'LSTM' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[13], line 86\u001B[0m\n\u001B[0;32m     81\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m record_dict\n\u001B[0;32m     84\u001B[0m epoch \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m20\u001B[39m\n\u001B[1;32m---> 86\u001B[0m model \u001B[38;5;241m=\u001B[39m \u001B[43mLSTM\u001B[49m()\n\u001B[0;32m     88\u001B[0m \u001B[38;5;66;03m# 1. 定义损失函数 采用交叉熵损失 (但是二分类)\u001B[39;00m\n\u001B[0;32m     89\u001B[0m loss_fct \u001B[38;5;241m=\u001B[39m F\u001B[38;5;241m.\u001B[39mbinary_cross_entropy_with_logits\n",
      "\u001B[1;31mNameError\u001B[0m: name 'LSTM' is not defined"
     ]
    }
   ],
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:28.107050Z",
     "start_time": "2025-03-19T05:50:28.107050Z"
    }
   },
   "source": [
    "#画线要注意的是损失是不一定在零到1之间的\n",
    "def plot_learning_curves(record_dict, sample_step=500):\n",
    "    # build DataFrame\n",
    "    train_df = pd.DataFrame(record_dict[\"train\"]).set_index(\"step\").iloc[::sample_step]\n",
    "    val_df = pd.DataFrame(record_dict[\"val\"]).set_index(\"step\")\n",
    "\n",
    "    # plot\n",
    "    fig_num = len(train_df.columns)\n",
    "    fig, axs = plt.subplots(1, fig_num, figsize=(5 * fig_num, 5))\n",
    "    for idx, item in enumerate(train_df.columns):    \n",
    "        axs[idx].plot(train_df.index, train_df[item], label=f\"train_{item}\")\n",
    "        axs[idx].plot(val_df.index, val_df[item], label=f\"val_{item}\")\n",
    "        axs[idx].grid()\n",
    "        axs[idx].legend()\n",
    "        # axs[idx].set_xticks(range(0, train_df.index[-1], 5000))\n",
    "        # axs[idx].set_xticklabels(map(lambda x: f\"{int(x/1000)}k\", range(0, train_df.index[-1], 5000)))\n",
    "        axs[idx].set_xlabel(\"step\")\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "plot_learning_curves(record, sample_step=10)  #横坐标是 steps"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 评估"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-19T05:50:28.108050Z",
     "start_time": "2025-03-19T05:50:28.107050Z"
    }
   },
   "source": [
    "# dataload for evaluating\n",
    "\n",
    "# load checkpoints\n",
    "model.load_state_dict(torch.load(\"checkpoints/imdb-lstm-subword/best.ckpt\", map_location=\"cpu\"))\n",
    "\n",
    "model.eval()\n",
    "loss, acc = evaluating(model, test_dl, loss_fct)\n",
    "print(f\"loss:     {loss:.4f}\\naccuracy: {acc:.4f}\")"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
